from time import sleep_us,ticks_us,sleep
from pyb import Pin,delay
import sensor, image, time, os, tf, utime, math
from pyb import UART
from pyb import Pin, Timer
import pyb
sensor.reset()
sensor.set_pixformat(sensor.RGB565)
sensor.set_framesize(sensor.QVGA)
sensor.set_windowing((240, 240))
sensor.skip_frames(time=2000)
net = "trained.tflite"
labels = [line.rstrip('\n') for line in open("labels.txt")]
labels_R = ['Recyclable_waste']
labels_K = ['Kitchen_waste']
labels_T = ['Harmful_waste']
labels_O = ['Other_waste']
uart = UART(3, 19200)
led_red = pyb.LED(1)
led_green = pyb.LED(2)
led_blue = pyb.LED(3)
light = Timer(2, freq=50000).channel(1, Timer.PWM, pin=Pin("P6"))
clock = time.clock()
def times_cal_2s():
	times = 0
	time_start = pyb.millis()
	print("识别结束，设备重启中...")
	while times <= 2000:
		time_End = pyb.millis()
		times = time_End - time_start
	print("重启成功，设备继续运作...")
def check_photo():
	times = 0
	time_start = pyb.millis()
	while times <= 1000:
		img = sensor.snapshot()
		time_End = pyb.millis()
		times = time_End - time_start
def garbage_Identification():
	light.pulse_width_percent(50)
	img1 = sensor.snapshot()
	for obj in tf.classify(net, img1, min_scale=1.0, scale_mul=0.8, x_overlap=0.5, y_overlap=0.5):
		predictions_list1 = list(zip(labels, obj.output()))
		max_index1 = predictions_list1.index(max(predictions_list1, key=lambda x: x[1]))
	max_index2 = max_index1
	predictions_list2 = predictions_list1
	while predictions_list2[max_index2][1] <= 0.9:
		img2 = sensor.snapshot()
		for obj in tf.classify(net, img2, min_scale=1.0, scale_mul=0.8, x_overlap=0.5, y_overlap=0.5):
			predictions_list2 = list(zip(labels, obj.output()))
			max_index2 = predictions_list2.index(max(predictions_list2, key=lambda x: x[1]))
	max_index3 = max_index2
	predictions_list3 = predictions_list2
	check_photo()
	while predictions_list3[max_index3][1]<= 0.9:
		img3 = sensor.snapshot()
		for obj in tf.classify(net, img3, min_scale=1.0, scale_mul=0.8, x_overlap=0.5, y_overlap=0.5):
			predictions_list3 = list(zip(labels, obj.output()))
			max_index3 = predictions_list3.index(max(predictions_list3, key=lambda x: x[1]))
	garbage = predictions_list3[max_index3][0]
	print('垃圾识别成功，输出识别结果......')
	print('垃圾名称：', garbage)
	garbage_name = garbage.strip()
	if garbage_name in labels_R:
		uart.write('r')
		print('垃圾类型：', '可回收垃圾')
	elif garbage_name in labels_K:
		uart.write('k')
		print('垃圾类型：', '厨余垃圾')
	elif garbage_name in labels_T:
		uart.write('t')
		print('垃圾类型：', '有害垃圾')
	elif garbage_name in labels_O:
		uart.write('o')
		print('垃圾类型：', '其它垃圾')
	light.pulse_width_percent(0)
print("设备初始化完成，成功开启......")
while True:
	light.pulse_width_percent(0)
	garbage_Identification()
	times_cal_2s()